#! /usr/bin/env python
#
#
import xarray as xr
import numpy as np
import matplotlib.pyplot as plt
import os
import urllib.request
import warnings

fluxes = ["rld", "rlu"]

def mae(diff, col_dim):
    #
    # Mean absolute error
    #
    return np.fabs(diff).mean(dim=col_dim)


def rms(diff, col_dim):
    #
    # Root mean square error
    #
    return np.sqrt(np.square(diff).mean(dim=col_dim))

def construct_lbl_esgf_root(var, esgf_node="llnl"):
    #
    # For a given variable name, provide the https URL for the LBLRM
    # line-by-line results
    #
    model = "LBLRTM-12-8"
    prefix = ("http://esgf3.dkrz.de/thredds/fileServer/cmip6/RFMIP/AER/" +
              model + "/rad-irf/r1i1p1f1/Efx/")
    if esgf_node == "llnl":
        prefix = ("http://esgf-data1.llnl.gov/thredds/fileServer/css03_data/"
                  "CMIP6/RFMIP/AER/" + model + "/rad-irf/r1i1p1f1/Efx/")
    return (prefix + var + "/gn/v20190514/" + var +
            "_Efx_" + model + "_rad-irf_r1i1p1f1_gn.nc")


######################
#######################

#
# Comparing reference and test results
#
if __name__ == '__main__':
    warnings.simplefilter("ignore", xr.SerializationWarning)

    lbl_suffix = "_Efx_LBLRTM-12-8_rad-irf_r1i1p1f1_gn.nc"
    gp_suffix  = "_Efx_RTE-RRTMGP-181204_rad-irf_r1i1p1f1_gn.nc"

    for v in fluxes:
        try:
            try:
                urllib.request.urlretrieve(construct_lbl_esgf_root(v),
                                           v + lbl_suffix)
            except:
                urllib.request.urlretrieve(
                    construct_lbl_esgf_root(v, esgf_node="dkrz"),
                    v + lbl_suffix)
        except:
            raise Exception("Failed to download {0}".format(v + lbl_suffix))

    lbl = xr.open_mfdataset([                                                                    v + lbl_suffix for v in fluxes], 
        combine="by_coords").sel(expt=0)
    old = xr.open_mfdataset([os.environ["RRTMGP_DATA"] + "/examples/rfmip-clear-sky/reference/"+ v + gp_suffix  for v in fluxes], 
        combine="by_coords").sel(expt=0)
    new = xr.open_mfdataset([                                      "examples/rfmip-clear-sky/" + v + gp_suffix  for v in fluxes], 
        combine="by_coords").sel(expt=0)

    #
    # Mean sbolute and RMS errors in upwelling flux
    #
    fig, axes = plt.subplots(ncols=2, figsize=[1.6 * 6, 6], sharey=True)

    for a, flux in enumerate(fluxes): 
        axes[a].plot(rms((old[flux]-lbl[flux]), col_dim="site"), lbl.level, c = "blue", label="Clough")
        axes[a].plot(rms((new[flux]-lbl[flux]), col_dim="site"), lbl.level, c = "red",  label="G-J-5")

        axes[a].plot(mae((old[flux]-lbl[flux]), col_dim="site"), lbl.level, "--", c = "blue")
        axes[a].plot(mae((new[flux]-lbl[flux]), col_dim="site"), lbl.level, "--", c = "red")
        axes[a].invert_yaxis()

        axes[a].set_title(flux)
        if(a == 1): axes[a].legend(frameon=False)
    
    fig.savefig("Flux-errors.png")

    #
    # Errors in net flux 
    #
    old_err_net = (old.rld - old.rlu) - (lbl.rld - lbl.rlu)
    new_err_net = (new.rld - new.rlu) - (lbl.rld - lbl.rlu)

    fig, axes = plt.subplots(ncols=2, figsize=[1.6 * 6, 6], sharey=True)
    axes[0].plot(rms(old_err_net, col_dim="site"), lbl.level, c = "blue", label="Clough")
    axes[0].plot(rms(new_err_net, col_dim="site"), lbl.level, c = "red", label="G-J-5")

    axes[0].plot(mae(old_err_net, col_dim="site"), lbl.level, "--", c = "blue")
    axes[0].plot(mae(new_err_net, col_dim="site"), lbl.level, "--", c = "red")

    axes[0].legend(frameon=False)
    axes[0].set_title("Net flux")

    #
    # Max error
    #
    axes[1].plot(np.abs(old.rld-lbl.rld).max(dim="site"), lbl.level, c = "blue")
    axes[1].plot(np.abs(new.rld-lbl.rld).max(dim="site"), lbl.level, c = "red")
    axes[1].plot(np.abs(old.rlu-lbl.rlu).max(dim="site"), lbl.level, c = "blue")
    axes[1].plot(np.abs(new.rlu-lbl.rlu).max(dim="site"), lbl.level, c = "red")

    axes[1].set_title("Max. abs. error")
    fig.savefig("Net-and-max-flux-error.png")